{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installation\n",
    "! pip install smolagents\n",
    "# To install from source instead of the last release, comment the command above and uncomment the following one.\n",
    "# ! pip install git+https://github.com/huggingface/smolagents.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📚 管理Agent的记忆"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "归根结底，Agent可以定义为由几个简单组件构成：它拥有工具、提示词。最重要的是，它具备对过往步骤的记忆，能够追溯完整的规划、执行和错误历史。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 回放Agent的记忆"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们提供了多项功能来审查Agent的过往运行记录。\n",
    "\n",
    "您可以通过插装（instrumentation）在可视化界面中查看Agent的运行过程，该界面支持对特定步骤进行缩放操作，具体方法参见[插装指南](https://huggingface.co/docs/smolagents/main/zh/tutorials/./inspect_runs)。\n",
    "\n",
    "您也可以使用`agent.replay()`方法实现回放：\n",
    "\n",
    "当Agent完成运行后："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smolagents import InferenceClientModel, CodeAgent\n",
    "\n",
    "agent = CodeAgent(tools=[], model=InferenceClientModel(), verbosity_level=0)\n",
    "\n",
    "result = agent.run(\"What's the 20th Fibonacci number?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "若要回放最近一次运行，只需使用："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent.replay()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 动态修改Agent的记忆"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "许多高级应用场景需要对Agent的记忆进行动态修改。\n",
    "\n",
    "您可以通过以下方式访问Agent的记忆："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smolagents import ActionStep\n",
    "\n",
    "system_prompt_step = agent.memory.system_prompt\n",
    "print(\"The system prompt given to the agent was:\")\n",
    "print(system_prompt_step.system_prompt)\n",
    "\n",
    "task_step = agent.memory.steps[0]\n",
    "print(\"\\n\\nThe first task step was:\")\n",
    "print(task_step.task)\n",
    "\n",
    "for step in agent.memory.steps:\n",
    "    if isinstance(step, ActionStep):\n",
    "        if step.error is not None:\n",
    "            print(f\"\\nStep {step.step_number} got this error:\\n{step.error}\\n\")\n",
    "        else:\n",
    "            print(f\"\\nStep {step.step_number} got these observations:\\n{step.observations}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用`agent.memory.get_full_steps()`可获取完整步骤字典数据。\n",
    "\n",
    "您还可以通过步骤回调（step callbacks）实现记忆的动态修改。\n",
    "\n",
    "步骤回调函数可通过参数直接访问`agent`对象，因此能够访问所有记忆步骤并根据需要进行修改。例如，假设您正在监控网页浏览Agent每个步骤的屏幕截图，希望保留最新截图同时删除旧步骤的图片以节省token消耗。\n",
    "\n",
    "可参考以下代码示例：\n",
    "_注：此代码片段不完整，部分导入语句和对象定义已精简，完整代码请访问[原始脚本](https://github.com/huggingface/smolagents/blob/main/src/smolagents/vision_web_browser.py)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import helium\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "from time import sleep\n",
    "\n",
    "def update_screenshot(memory_step: ActionStep, agent: CodeAgent) -> None:\n",
    "    sleep(1.0)  # Let JavaScript animations happen before taking the screenshot\n",
    "    driver = helium.get_driver()\n",
    "    latest_step = memory_step.step_number\n",
    "    for previous_memory_step in agent.memory.steps:  # Remove previous screenshots from logs for lean processing\n",
    "        if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number <= latest_step - 2:\n",
    "            previous_memory_step.observations_images = None\n",
    "    png_bytes = driver.get_screenshot_as_png()\n",
    "    image = Image.open(BytesIO(png_bytes))\n",
    "    memory_step.observations_images = [image.copy()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后在初始化Agent时，将此函数传入`step_callbacks`参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CodeAgent(\n",
    "    tools=[WebSearchTool(), go_back, close_popups, search_item_ctrl_f],\n",
    "    model=model,\n",
    "    additional_authorized_imports=[\"helium\"],\n",
    "    step_callbacks=[update_screenshot],\n",
    "    max_steps=20,\n",
    "    verbosity_level=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "请访问我们的 [vision web browser code](https://github.com/huggingface/smolagents/blob/main/src/smolagents/vision_web_browser.py) 查看完整可运行示例。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 分步运行 Agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当您需要处理耗时数天的工具调用时，这种方式特别有用：您可以逐步执行Agents。这还允许您在每一步更新记忆。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smolagents import InferenceClientModel, CodeAgent, ActionStep, TaskStep\n",
    "\n",
    "agent = CodeAgent(tools=[], model=InferenceClientModel(), verbosity_level=1)\n",
    "print(agent.memory.system_prompt)\n",
    "\n",
    "task = \"What is the 20th Fibonacci number?\"\n",
    "\n",
    "# You could modify the memory as needed here by inputting the memory of another agent.\n",
    "# agent.memory.steps = previous_agent.memory.steps\n",
    "\n",
    "# Let's start a new task!\n",
    "agent.memory.steps.append(TaskStep(task=task, task_images=[]))\n",
    "\n",
    "final_answer = None\n",
    "step_number = 1\n",
    "while final_answer is None and step_number <= 10:\n",
    "    memory_step = ActionStep(\n",
    "        step_number=step_number,\n",
    "        observations_images=[],\n",
    "    )\n",
    "    # Run one step.\n",
    "    final_answer = agent.step(memory_step)\n",
    "    agent.memory.steps.append(memory_step)\n",
    "    step_number += 1\n",
    "\n",
    "    # Change the memory as you please!\n",
    "    # For instance to update the latest step:\n",
    "    # agent.memory.steps[-1] = ...\n",
    "\n",
    "print(\"The final answer is:\", final_answer)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
